import polars as pl

from api import sqlpyodbc,sql_polars
from pydantic import BaseModel
from typing import Optional


class PackageMaterial(BaseModel):
    code: str
    type: str
    description: Optional[str] = None
    unit: str
    operator: Optional[str] = None

def get_materials(search_text: str = '', page: int = 1, page_size: int = 10):
    # 首先获取总记录数
    count_sql = """
        SELECT COUNT(*) as total
        FROM PackageMaterial 
        WHERE IsDeleted = 0
    """
    if search_text:
        count_sql += f" AND (Code LIKE '%{search_text}%' OR Type LIKE '%{search_text}%')"
    
    try:
        total_result = sql_polars.query(count_sql)
        total = int(total_result.select("total")[0, 0]) if total_result.height > 0 else 0

        # 计算分页
        offset = (page - 1) * page_size
        
        # 获取分页数据
        sql = """
            SELECT Code as code, Type as type, Description as description, 
                   Unit as unit, CreateTime as createTime, CreateBy as creator,
                   ISNULL(CurrentStock, 0) as currentStock
            FROM PackageMaterial 
            WHERE IsDeleted = 0
        """
        if search_text:
            sql += f" AND (Code LIKE '%{search_text}%' OR Type LIKE '%{search_text}%')"
        
        sql += f" ORDER BY CreateTime DESC OFFSET {offset} ROWS FETCH NEXT {page_size} ROWS ONLY"
        
        result = sql_polars.query(sql)
        result = result.with_columns(
            pl.col("createTime").dt.strftime("%Y-%m-%d %H:%M:%S").alias("createTime")
        )
        if result is not None:
            records = result.to_dicts()
            return {
                "records": records,
                "total": total
            }
        return None
    except Exception as e:
        print(f"Error in get_materials: {str(e)}")
        return None

def add_material(material: PackageMaterial):
    sql = f"""
        INSERT INTO PackageMaterial (Code, Type, Description, Unit, CreateBy)
        VALUES ('{material.code}', '{material.type}', 
                '{material.description}', '{material.unit}', '{material.operator or "system"}')
    """
    try:
        return sqlpyodbc.execute(sql)
    except:
        return None

def update_material(material: PackageMaterial):
    sql = f"""
        UPDATE PackageMaterial 
        SET Type = '{material.type}',
            Description = '{material.description}',
            Unit = '{material.unit}',
            UpdateTime = GETDATE(),
            UpdateBy = '{material.operator or "system"}'
        WHERE Code = '{material.code}'
    """
    try:
        return sqlpyodbc.execute(sql)
    except:
        return None

def delete_material(code: str, operator: str = 'system'):
    sql = f"""
        UPDATE PackageMaterial 
        SET IsDeleted = 1,
            UpdateTime = GETDATE(),
            UpdateBy = '{operator}'
        WHERE Code = '{code}'
    """
    try:
        return sqlpyodbc.execute(sql)
    except:
        return None 